use crate::connection::{Connection, ConnectionManager};

use crate::error::{Result, WsError};
use crate::protocol::Message;
use futures_util::{SinkExt, StreamExt};
use std::sync::Arc;
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio::time::{timeout, Duration};
use tokio_tungstenite::{accept_async};
use tracing::{debug, error, info, warn};
use tungstenite::Message as WsMessage;
use uuid::Uuid;

pub struct WebSocketHandler {
    connection_manager: Arc<ConnectionManager>,
}

impl WebSocketHandler {
    pub fn new(connection_manager: Arc<ConnectionManager>) -> Self {
        Self { connection_manager }
    }
    
    pub async fn handle_connection(&self, stream: TcpStream, addr: std::net::SocketAddr) -> Result<()> {
        info!("New WebSocket connection from: {}", addr);
        
        let ws_stream = accept_async(stream).await?;
        let (mut ws_sender, mut ws_receiver) = ws_stream.split();
        
        // Create channel for sending messages to this connection
        let (tx, mut rx) = mpsc::unbounded_channel::<WsMessage>();
        
        // Generate unique connection ID
        let connection_id = Uuid::new_v4();
        
        // Wait for initial join message to get username and user_id
        let (username, user_id) = match timeout(Duration::from_secs(10), ws_receiver.next()).await {
            Ok(Some(Ok(WsMessage::Text(text)))) => {
                match Message::from_json(&text) {
                    Ok(Message::Join { user_id, username }) => {
                        info!("User {} joined with connection ID: {}", username, connection_id);
                        (username, user_id)
                    }
                    _ => {
                        warn!("Invalid initial message from {}", addr);
                        return Err(WsError::InvalidMessageFormat);
                    }
                }
            }
            _ => {
                warn!("No valid join message received from {} within timeout", addr);
                return Err(WsError::InvalidMessageFormat);
            }
        };
        
        // Create connection
        let connection = Connection::new(connection_id, username.clone(), tx);
        self.connection_manager.add_connection(connection, user_id);
        
        // Send join confirmation
        let join_ack = Message::Join {
            user_id: user_id,
            username: username.clone(),
        };
        
        if let Ok(json) = join_ack.to_json() {
            if let Err(e) = ws_sender.send(WsMessage::Text(json)).await {
                error!("Failed to send join acknowledgment: {}", e);
                self.connection_manager.remove_connection(&connection_id);
                return Err(e.into());
            }
        }
        
        // Send current online users to the new connection
        let online_users = self.connection_manager.get_all_connections();
        for user_info in online_users {
            if user_info.id != connection_id {
                // Get the actual user ID for this connection
                if let Some(existing_user_id) = self.connection_manager.get_user_id_by_connection_id(&user_info.id) {
                    let user_join_msg = Message::Join {
                        user_id: existing_user_id,
                        username: user_info.username.clone(),
                    };
                    if let Ok(json) = user_join_msg.to_json() {
                        if let Err(e) = ws_sender.send(WsMessage::Text(json)).await {
                            error!("Failed to send existing user info to new connection: {}", e);
                        }
                    }
                }
            }
        }
        
        // Broadcast join message to all other connections
        let join_msg = Message::Join {
            user_id: user_id,
            username: username.clone(),
        };
        
        if let Err(e) = self.connection_manager.broadcast_message(&join_msg, Some(connection_id)).await {
            error!("Failed to broadcast join message: {}", e);
        }
        
        let connection_id_copy = connection_id;
        
        // Spawn task to handle outgoing messages
        let outgoing_task = tokio::spawn(async move {
            let mut ws_sender = ws_sender;
            
            while let Some(message) = rx.recv().await {
                if let Err(e) = ws_sender.send(message).await {
                    error!("Failed to send message to connection {}: {}", connection_id_copy, e);
                    break;
                }
            }
            
            info!("Outgoing message handler stopped for connection {}", connection_id_copy);
        });
        
        let connection_manager_clone = Arc::clone(&self.connection_manager);
        let connection_id_clone = connection_id;
        
        // Handle incoming messages
        let result = async {
            while let Some(msg) = ws_receiver.next().await {
                match msg {
                    Ok(WsMessage::Frame(_)) => {
                        // Ignore raw frames
                        continue;
                    }
                    Ok(WsMessage::Text(text)) => {
                        connection_manager_clone.increment_received();
                        
                        match Message::from_json(&text) {
                            Ok(message) => {
                                if let Err(e) = self.handle_message(connection_id_clone, message).await {
                                    error!("Error handling message from {}: {}", connection_id_clone, e);
                                }
                            }
                            Err(e) => {
                                warn!("Invalid message format from {}: {}", connection_id_clone, e);
                            }
                        }
                    }
                    Ok(WsMessage::Binary(_)) => {
                        warn!("Binary messages not supported from {}", connection_id_clone);
                    }
                    Ok(WsMessage::Ping(_)) => {
                        // Pong will be sent automatically by tungstenite
                        debug!("Ping received from {}", connection_id_clone);
                    }
                    Ok(WsMessage::Pong(_)) => {
                        debug!("Pong received from {}", connection_id_clone);
                    }
                    Ok(WsMessage::Close(_)) => {
                        info!("Connection {} requested close", connection_id_clone);
                        break;
                    }
                    Err(e) => {
                        error!("WebSocket error from {}: {}", connection_id_clone, e);
                        break;
                    }
                }
            }
            
            Ok::<(), WsError>(())
        }.await;
        
        // Clean up connection
        info!("Connection {} disconnected", connection_id);
        self.connection_manager.remove_connection(&connection_id);
        
        // Broadcast leave message
        let leave_msg = Message::Leave {
            user_id: connection_id,
        };
        
        if let Err(e) = self.connection_manager.broadcast_message(&leave_msg, None).await {
            error!("Failed to broadcast leave message: {}", e);
        }
        
        // Wait for outgoing task to complete
        let _ = outgoing_task.await;
        
        result
    }
    
    async fn handle_message(&self, sender_id: Uuid, message: Message) -> Result<()> {
        debug!("Handling message from {}: {:?}", sender_id, message);
        
        match message {
            Message::Broadcast { .. } => {
                self.connection_manager.broadcast_message(&message, Some(sender_id)).await?;
            }
            Message::Direct { target_id, .. } => {
                self.connection_manager.send_direct_message(target_id, &message).await?;
            }
            Message::Ping => {
                let pong = Message::Pong;
                self.connection_manager.send_direct_message(sender_id, &pong).await?;
            }
            Message::Pong => {
                // Just update activity timestamp
                if let Some(conn) = self.connection_manager.get_connection(&sender_id) {
                    conn.update_activity();
                }
            }
            _ => {
                warn!("Unexpected message type from {}: {:?}", sender_id, message);
            }
        }
        
        Ok(())
    }
}